#得到TaxID列表后的数据清洗并获取分类阶元等信息
import pandas as pd
# 读取TaxID与分类阶元信息的映射文件
def read_lineage(file_path):
    lineage_dict = {}
    with open(file_path, 'r') as f:
        for line in f:
            parts = line.strip().split('\t')
            if len(parts) == 2:
                taxid, lineage = parts
                lineage_dict[taxid] = lineage
            else:
                print(f"Warning: Line format is incorrect in {file_path}: {line.strip()}")
    return lineage_dict
# 读取新生成的文件并提取指定列
def read_and_extract_data(file_path, lineage_dict):
    df = pd.read_csv(file_path, sep='\t')
    # 确保列名存在
    required_columns = ['assembly_accession', 'assembly_level', 'genome_size', 'protein_coding_gene_count', 'group', 'organism_name', 'species_taxid']
    if not all(col in df.columns for col in required_columns):
        print(f"Error: Some required columns are missing in {file_path}. Available columns: {df.columns.tolist()}")
        return pd.DataFrame()
    # 提取指定列
    extracted_df = df[required_columns]
    # 确保species_taxid列是字符串类型
    extracted_df['species_taxid'] = extracted_df['species_taxid'].astype(str)
    # 添加分类阶元信息
    extracted_df['lineage'] = extracted_df['species_taxid'].map(lineage_dict)
    # 检查是否有未匹配的行
    unmatched_count = extracted_df['lineage'].isna().sum()
    if unmatched_count > 0:
        print(f"Warning: {unmatched_count} entries have no matching lineage information.")
    return extracted_df
# 处理lineage列，提取目、科、种信息，并将其他阶元信息动态拆分为多个列
def process_lineage(lineage, organism_name):
    if pd.isna(lineage):
        return pd.Series([None, None, None, organism_name], index=['order', 'family', 'genus', 'species'])
    parts = lineage.split(';')
    order = family = genus = None
    # 提取目、科、属信息
    for part in parts:
        part = part.strip()
        if part.endswith('dae'):
            family = part
        elif part.endswith(('era', 'ida', 'aca', 'mes', 'nes', 'ata', 'eae', 'tia', 'nia', 'hia', 'yla', 'ota', 'ora', 'ona', 'ria')):
            order = part
        elif genus is None:
            genus = part
    # 将剩余的分类阶元信息动态拆分为多个列，动态生成列名
    other_taxa = [part.strip() for part in parts if part not in [order, family, genus]]
    other_taxa_columns = [f'taxon_{i}' for i in range(len(other_taxa))]
    species = organism_name
    return pd.Series([order, family, genus, species] + other_taxa, index=['order', 'family', 'genus', 'species'] + other_taxa_columns)

if __name__ == "__main__":
    # 读取分类阶元信息
    lineage_dict = read_lineage('lineage.txt')
    print(f"Loaded {len(lineage_dict)} lineage entries.")
    # 读取新生成的文件并提取指定列
    final_df = read_and_extract_data('merged_assembly_summary.txt', lineage_dict)
    if not final_df.empty:
        # 处理lineage列，提取目、科、种信息，并将其他阶元信息动态拆分为多个列
        processed_data = final_df.apply(lambda row: process_lineage(row['lineage'], row['organism_name']), axis=1)
        # 将处理后的数据合并回主DataFrame
        final_df = pd.concat([final_df.drop(columns=['lineage']), processed_data], axis=1)
        # 输出结果
        final_df.to_csv('final_assembly_summary_with_lineage.csv', index=False)
        print(f"Saved final data to final_assembly_summary_with_lineage.csv with {len(final_df)} entries.")
    else:
        print("No data to process. Exiting.")